from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 配置基础数据 ==========
VISIT_TYPES = ['门诊', '住院', '急诊']
DIAG_TYPES = ['初步诊断', '确诊', '复诊', '补充诊断']
CONFIRMED_MARKS = ['是', '否']
MAINDIAG_MARKS = ['是', '否']
COMPLICATION_MARKS = ['是', '否']
DEPARTMENTS = ['肿瘤科', '胸外科', '普外科', '乳腺外科', '消化科']

# 癌症相关诊断
CANCER_DIAGNOSES = [
    ("C34.1", "支气管癌"),
    ("C34.9", "肺癌"),
    ("C50.9", "乳腺癌"),
    ("C18.9", "结直肠癌"),
    ("C16.9", "胃癌"),
    ("C15.9", "食管癌"),
    ("C22.0", "肝癌"),
    ("C25.9", "胰腺癌"),
]


# ========== 时间范围配置 ==========
def generate_DAY_RANGE(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    days = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y-%m-%d"))
        current += relativedelta(days=1)
    return days


DAY_RANGE = generate_DAY_RANGE("2025-08-01", "2025-08-31")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 150 if date_obj.day % 2 == 0 else 120


# ========== 核心生成函数 ==========
def generate_diagnosis_records():
    sql_statements = []
    record_count = 0

    patient_id_counter = 2001  # 和 b02_1 区分开
    visit_sn_counter = 3001
    diag_id_counter = 4001

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            record_count += 1

            patient_id = f"PT{patient_id_counter}"
            patient_id_counter += 1

            visit_sn = f"VIS{visit_sn_counter}"
            visit_sn_counter += 1

            diag_id = f"DIA{diag_id_counter}"
            diag_id_counter += 1

            diag_code, diag_name = random.choice(CANCER_DIAGNOSES)
            diag_datetime = fake.date_time_between(
                start_date=datetime.strptime(day, "%Y-%m-%d"),
                end_date=(datetime.strptime(day, "%Y-%m-%d") + timedelta(days=1))
            ).strftime('%Y-%m-%d %H:%M:%S')

            data = {
                'patient_id': patient_id,
                'visit_sn': visit_sn,
                'visit_type': random.choice(VISIT_TYPES),
                'medical_record_no': f"MR{random.randint(100000, 999999)}",
                'outpatient_no': f"OUT{random.randint(1000, 9999)}",
                'visit_times': random.randint(1, 5),
                'inpatient_no': f"IP{random.randint(100000, 999999)}",
                'hospitalization_times': str(random.randint(1, 3)),
                'name': fake.name(),
                'diag_id': diag_id,
                'diag_serial_number': str(random.randint(1, 10)),
                'diag_type': random.choice(DIAG_TYPES),
                'diag_code': diag_code,
                'diag_name': diag_name,
                'diag_explanation': f"{diag_name} 临床相关说明",
                'diag_datetime': diag_datetime,
                'confirmed_diag_mark': random.choice(CONFIRMED_MARKS),
                'maindiag_mark': random.choice(MAINDIAG_MARKS),
                'diag_doctor_no': f"DOC{random.randint(100, 999)}",
                'is_complication': random.choice(COMPLICATION_MARKS),
                'create_diagnose_dept_name': random.choice(DEPARTMENTS),
                'record_status': random.randint(0, 1),
                'yy_etl_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'yy_upload_status': 0,
                'yy_upload_time': None,
                'yy_batch_time': day,  # 按天作为批次
                'yy_record_batch_id': f"BATCH{day}_{i}",
                'yy_backfill_time': None,
                'yy_backfill_status': None
            }

            sql = _generate_sql('b02_3', data)
            sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_diagnosis_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b02_3')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "b02_3_diagnosis_records.sql")
